//
// Create by RangiLyu
// 2020 / 10 / 2
//

#ifndef NANODET_H
#define NANODET_H

#include "net.h"
#include <opencv2/core/core.hpp>

typedef struct HeadInfo_
{
    std::string cls_layer;
    std::string dis_layer;
    int stride;
} HeadInfo;

typedef struct CenterPrior_
{
    int x;
    int y;
    int stride;
} CenterPrior;

typedef struct BoxInfo {
    float x1;
    float y1;
    float x2;
    float y2;
    float score;
    int label;
} BoxInfo;



class NanoDet{
public:

    NanoDet(AAssetManager *mgr, const char *param, const char *bin, bool useGPU);

    ~NanoDet();
    std::vector<BoxInfo> detect_bitmap(JNIEnv *env, jobject image, float score_threshold, float nms_threshold);
    int detect(const cv::Mat& rgb,std::vector<BoxInfo>& dets, float score_threshold = 0.7f, float nms_threshold = 0.7f);
    int draw(const cv::Mat& rgb, const std::vector<BoxInfo>& objects);
    std::vector<std::string> labels{"0","1"};
private:
    void decode_infer(ncnn::Mat& feats, std::vector<CenterPrior>& center_priors, float threshold, std::vector<std::vector<BoxInfo>>& results, float width_ratio, float height_ratio);
    BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride, float width_ratio, float height_ratio);

    static void nms(std::vector<BoxInfo>& result, float nms_threshold);

    ncnn::Net *Net;
    // modify these parameters to the same with your config if you want to use your own model
    int input_size[2] = {320, 320}; // input height and width
    int num_class = 2; // number of classes. 80 for COCO
    int reg_max = 7; // `reg_max` set in the training config. Default: 7.
    std::vector<int> strides = { 8, 16, 32, 64 }; // strides of the multi-level feature.


public:
    static NanoDet *detector;
    static bool hasGPU;
};


#endif //NANODET_H
